package Postgre;

import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

import java.sql.*;

import static org.testng.Assert.assertEquals;



/**
 * Created by IntelliJ IDEA.
 * User: Andrey
 * Date: 16.03.12
 * Time: 20:11
 * To change this template use File | Settings | File Templates.
 */
public class testGetPrimaryKey {
    private Connection con;

    @BeforeClass
    public void connect() throws ClassNotFoundException, SQLException {
        Class.forName("org.postgresql.Driver");
        con = DriverManager.getConnection("jdbc:postgresql://localhost:5432/test", "andrey", "andrey");

        Statement st = con.createStatement();
        st.executeUpdate("create table test1 (id integer, value varchar(80), primary key(id))");
        st.executeUpdate("create table test2 (id integer, value integer, primary key(value))");
    }

    @AfterClass
    public void closeConnection() throws SQLException {
        Statement st = con.createStatement();
        st.execute("drop table test1");
        st.execute("drop table test2");

        con.close();
    }

    @Test
    public void testGetTables() throws SQLException {
        String pKey;
        Statement st = con.createStatement();

        ResultSet rs = st.executeQuery("SELECT c2.relname, a.attname \n" +
                "FROM pg_class c, pg_class c2, pg_index i, pg_attribute a \n" +
                "WHERE c.relname = 'test1' AND c.oid = i.indrelid AND i.indexrelid = c2.oid \n" +
                "AND i.indisprimary AND i.indisunique \n" +
                "AND a.attrelid=c2.oid \n" +
                "AND a.attnum>0; ");

        rs.next();
        pKey = rs.getString("attname");
        System.err.println("Primary key: " + pKey);
        assertEquals(pKey, "id");


        rs = st.executeQuery("SELECT c2.relname, a.attname \n" +
                "FROM pg_class c, pg_class c2, pg_index i, pg_attribute a \n" +
                "WHERE c.relname = 'test2' AND c.oid = i.indrelid AND i.indexrelid = c2.oid \n" +
                "AND i.indisprimary AND i.indisunique \n" +
                "AND a.attrelid=c2.oid \n" +
                "AND a.attnum>0; ");
        rs.next();
        pKey = rs.getString("attname");
        System.err.println("Primary key: " + pKey);
        assertEquals(pKey, "value");
    }
}
